import io
from collections.abc import Iterable

import cv2
import numpy as np
from .preprocessing import Preprocessing
import av

class StreamPreprocessor(Preprocessing):
    def preprocess(self, stream_bytes: bytes) -> np.ndarray:
        """
        预处理视频流（以H.264 over RTP或其他类似方式编码）以供Triton推理使用。

        参数:
            stream_bytes: 编码后的视频流数据。

        返回值:
            np.ndarray: 预处理后的视频帧，以numpy数组形式返回，数组形状为(num_frames, height, width, channels)。
        """
        try:
            # 假设有一个能够处理传入流格式的解码器
            decoded_frames = decode_video_stream(stream_bytes)
        except Exception as e:
            # 对于解码过程出现的任何错误，进行适当的异常处理
            # 这里举例打印错误信息并返回一个空数组，实际应用中可能需要更复杂的错误处理逻辑
            print(f"Error decoding video stream: {e}")
            return np.empty((0, 640, 640, 3), dtype=np.uint8)

        frames = []
        for frame in decoded_frames:
            try:
                # 执行必要的帧调整大小、标准化等操作
                # 例如，调整到固定大小并转换为BGR格式：
                resized_frame = cv2.resize(frame, (640, 640))
                frames.append(resized_frame)
            except Exception as e:
                # 对于处理单个帧时出现的错误，进行适当的异常处理
                # 这里举例跳过当前帧，实际应用中可能需要根据具体情况进行处理
                print(f"Error processing frame: {e}")
                continue

        # 分批处理帧以减少内存占用，假设每50帧处理一次
        batch_size = 50
        for i in range(0, len(frames), batch_size):
            batch_frames = frames[i:i+batch_size]
            if batch_frames:
                yield np.stack(batch_frames)

        # 如果需要将所有帧一次性返回，而非使用生成器，可以取消以下注释
        # return np.stack(frames)


def decode_video_stream(stream_bytes: bytes) -> Iterable[np.ndarray]:
    """
    解码视频流（以H.264 over RTP或其他类似方式编码）。

    参数:
        stream_bytes: 编码后的视频流数据。

    返回值:
        Iterable[np.ndarray]: 解码后的视频帧迭代器，每个元素为一个numpy数组表示的视频帧。
    """
    container = av.open(io.BytesIO(stream_bytes))

    for packet in container.demux():
        if packet.stream.type == 'video':
            for frame in packet.decode():
                # 将解码后的帧转换为 numpy 数组
                yield frame.to_ndarray(format='bgr24')
